package com.yupi.aicodehelper.ai.store;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.yupi.aicodehelper.ai.AiCodeHelper;
import com.yupi.aicodehelper.ai.dto.ChatSession;
import com.yupi.aicodehelper.ai.dto.ChatSessionMapper;
import com.yupi.aicodehelper.ai.dto.StoredMessage;
import com.yupi.aicodehelper.ai.dto.TimedMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.Jedis;
import dev.langchain4j.data.message.*;

import javax.sql.DataSource;
import java.sql.*;
import java.time.Instant;
import java.util.*;
import java.util.Date;
import java.util.concurrent.*;
import java.util.stream.Collectors;

// Redis + MySQL 持久化 ChatMemoryStore
@Slf4j
public class RedisChatMemoryStore implements ChatMemoryStore {

    private static final String REDIS_PREFIX = "chat_memory:";
    private static final int EXPIRE_SECONDS = 24 * 60 * 60; // 1天
    private static final int MAX_MESSAGES = 200; // 最多保留200条，防止key过大

    private final JedisPool jedisPool;
    private final ObjectMapper objectMapper = new ObjectMapper();

    // 注入大模型调用 Service
    private final AiCodeHelper aiCodeHelper;
    // 注入 Mapper/Repository
    private final ChatSessionMapper chatSessionMapper;

    public RedisChatMemoryStore(JedisPool jedisPool, AiCodeHelper aiCodeHelper, ChatSessionMapper chatSessionMapper) {
        this.jedisPool = jedisPool;
        this.aiCodeHelper = aiCodeHelper;
        this.chatSessionMapper = chatSessionMapper;
    }

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        String key = REDIS_PREFIX + memoryId;
        long start = System.currentTimeMillis();
        try (Jedis jedis = jedisPool.getResource()) {
            String json = jedis.get(key);
            long cost = System.currentTimeMillis() - start;
            log.info("[Redis GET] memoryId={} cost={}ms", memoryId, cost);

            if (json != null && !json.isEmpty()) {
                List<StoredMessage> list = objectMapper.readValue(json, new TypeReference<List<StoredMessage>>() {});
                List<ChatMessage> messages = new ArrayList<>();
                for (StoredMessage m : list) {
                    messages.add(m.toChatMessage());
                }
                return messages;
            }
        } catch (Exception e) {
            long cost = System.currentTimeMillis() - start;
            log.error("[Redis GET] 读取 Redis 出错 memoryId={} cost={}ms", memoryId, cost, e);
        }

        return Collections.emptyList();
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        String key = REDIS_PREFIX + memoryId;
        long start = System.currentTimeMillis();

        List<StoredMessage> storedList = new ArrayList<>();
        long now = System.currentTimeMillis();

        for (ChatMessage msg : messages) {
            StoredMessage sm = StoredMessage.fromChatMessage((int) memoryId, msg);
            if (sm.getCreatedAt() == 0) sm.setCreatedAt(now);
            storedList.add(sm);
        }

        // 限制消息条数
        if (storedList.size() > MAX_MESSAGES) {
            storedList = storedList.subList(storedList.size() - MAX_MESSAGES, storedList.size());
        }

        try (Jedis jedis = jedisPool.getResource()) {
            boolean isNewSession = !jedis.exists(key); // 判断是否新对话
            String json = objectMapper.writeValueAsString(storedList);
            jedis.setex(key, EXPIRE_SECONDS, json);

            long cost = System.currentTimeMillis() - start;
            log.info("[Redis SET] memoryId={} messages={} cost={}ms", memoryId, storedList.size(), cost);

            // 如果是新会话，异步生成标题并落表
            if (isNewSession && !messages.isEmpty()) {
                String firstUserQuestion = storedList.get(0).getText();
                CompletableFuture.runAsync(() -> generateAndSaveTitle(memoryId.toString(), firstUserQuestion));
            }
        } catch (Exception e) {
            long cost = System.currentTimeMillis() - start;
            log.error("[Redis SET] 写入 Redis 出错 memoryId={} cost={}ms", memoryId, cost, e);
        }
    }

    @Override
    public void deleteMessages(Object memoryId) {
        String key = REDIS_PREFIX + memoryId;
        long start = System.currentTimeMillis();
        try (Jedis jedis = jedisPool.getResource()) {
            jedis.del(key);
            long cost = System.currentTimeMillis() - start;
            log.info("[Redis DEL] memoryId={} cost={}ms", memoryId, cost);
        } catch (Exception e) {
            long cost = System.currentTimeMillis() - start;
            log.error("[Redis DEL] 删除 Redis 出错 memoryId={} cost={}ms", memoryId, cost, e);
        }
    }

    /**
     * 异步调用大模型生成标题并入库
     */
    private void generateAndSaveTitle(String memoryId, String userQuestion) {
        try {
            String prompt = "请从以下用户问题中提炼一个简短的对话标题，不超过15个字，简洁明了：\n" + userQuestion;
            String title = aiCodeHelper.chat(prompt);

            ChatSession session = new ChatSession();
            session.setMemoryId(memoryId);
            session.setTitle(title);
            session.setCreatedAt(new Date());

            chatSessionMapper.insert(session);

            log.info("[ChatSession] 新会话已创建 memoryId={} title={}", memoryId, title);
        } catch (Exception e) {
            log.error("[ChatSession] 创建会话标题失败 memoryId={} userQuestion={}", memoryId, userQuestion, e);
        }
    }
}
